package top.zhangjianyong.blog.utils;

import com.maxmind.geoip2.DatabaseReader;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

/**
 * 通用 DatabaseReader 连接池
 * 支持任意类型的 DatabaseReader 实例管理
 * 线程安全，使用阻塞队列提高性能
 */
@Slf4j
public class DatabaseReaderPool {

    /**
     * -- GETTER --
     * 获取池大小
     */
    // 连接池配置
    @Getter
    private final int poolSize;
    private final int waitTimeout;
    /**
     * -- GETTER --
     * 获取连接池名称
     */
    @Getter
    private final String poolName;

    // 读取器池 - 使用阻塞队列
    private final BlockingQueue<DatabaseReader> readerPool = new LinkedBlockingQueue<>();
    private final AtomicInteger readerCount = new AtomicInteger(0);
    private final Object poolLock = new Object();

    // 创建读取器的工厂方法
    private final Supplier<DatabaseReader> readerFactory;

    /**
     * 构造函数
     *
     * @param poolName      连接池名称（如 "ASN" 或 "City"）
     * @param poolSize      连接池大小
     * @param waitTimeout   等待超时时间（毫秒）
     * @param readerFactory 创建读取器的工厂方法
     */
    public DatabaseReaderPool(String poolName, int poolSize, int waitTimeout, Supplier<DatabaseReader> readerFactory) {
        this.poolName = poolName;
        this.poolSize = poolSize;
        this.waitTimeout = waitTimeout;
        this.readerFactory = readerFactory;
        log.info("初始化 {} DatabaseReader 连接池，池大小: {}", poolName, poolSize);
    }

    /**
     * 获取 DatabaseReader 实例
     *
     * @return DatabaseReader 实例，如果获取失败返回 null
     */
    public DatabaseReader getReader() {
        try {
            // 首先尝试从池中获取（非阻塞）
            DatabaseReader reader = readerPool.poll();
            if (reader != null) {
                return reader;
            }

            // 池中没有可用实例，检查是否可以创建新的
            if (readerCount.get() < poolSize) {
                synchronized (poolLock) {
                    if (readerCount.get() < poolSize) {
                        reader = createReader();
                        if (reader != null) {
                            readerCount.incrementAndGet();
                            log.info("创建新的 {} DatabaseReader，当前实例数: {}", poolName, readerCount.get());
                            return reader;
                        }
                    }
                }
            }

            // 如果无法创建新实例，说明池已满，此时使用阻塞等待
            // 因为此时池中应该有实例，只是都被占用了
            reader = readerPool.poll(waitTimeout, TimeUnit.MILLISECONDS);
            if (reader != null) {
                return reader;
            }

            log.warn("无法获取 {} DatabaseReader，池中无可用实例且已达到最大数量限制", poolName);
            return null;

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("获取 {} DatabaseReader 时被中断", poolName);
            return null;
        }
    }

    /**
     * 归还 DatabaseReader 到池中
     *
     * @param reader DatabaseReader 实例
     */
    public void returnReader(DatabaseReader reader) {
        if (reader != null) {
            // 使用 offer 方法，如果队列满了会返回 false
            boolean offered = readerPool.offer(reader);
            if (!offered) {
                log.warn("{} 连接池已满，无法归还 DatabaseReader 实例", poolName);
                // 如果池满了，直接关闭这个实例
                try {
                    reader.close();
                } catch (Exception e) {
                    log.error("关闭 {} DatabaseReader 失败", poolName, e);
                }
            }
        }
    }

    /**
     * 创建 DatabaseReader 实例
     *
     * @return DatabaseReader 实例，创建失败返回 null
     */
    private DatabaseReader createReader() {
        try {
            return readerFactory.get();
        } catch (Exception e) {
            log.error("创建 {} DatabaseReader 失败", poolName, e);
            return null;
        }
    }

    /**
     * 获取连接池状态信息
     *
     * @return 连接池状态字符串
     */
    public String getPoolStatus() {
        return String.format("%s Pool: %d/%d (可用: %d)",
            poolName, readerCount.get(), poolSize, readerPool.size());
    }

    /**
     * 获取当前实例数
     *
     * @return 当前实例数
     */
    public int getCurrentCount() {
        return readerCount.get();
    }

    /**
     * 获取可用实例数
     *
     * @return 可用实例数
     */
    public int getAvailableCount() {
        return readerPool.size();
    }

    /**
     * 关闭连接池，释放所有资源
     */
    public void shutdown() {
        log.info("关闭 {} DatabaseReader 连接池", poolName);

        // 关闭所有读取器
        DatabaseReader reader;
        while ((reader = readerPool.poll()) != null) {
            try {
                reader.close();
            } catch (Exception e) {
                log.error("关闭 {} DatabaseReader 失败", poolName, e);
            }
        }

        readerCount.set(0);

        log.info("{} DatabaseReader 连接池已关闭", poolName);
    }
} 